package com.btl.service.ai.support.rag.embedding.models;

import dev.langchain4j.model.embedding.onnx.AbstractInProcessEmbeddingModel;
import dev.langchain4j.model.embedding.onnx.OnnxBertBiEncoder;
import dev.langchain4j.model.embedding.onnx.PoolingMode;

import java.nio.file.Path;
import java.util.concurrent.Executor;

import static dev.langchain4j.internal.ValidationUtils.ensureNotNull;

public class BegSmallZhV15Model extends AbstractInProcessEmbeddingModel {


    private final OnnxBertBiEncoder MODEL;

    /**
     * Creates an instance of an {@code BegSmallZhV15Model}.
     * Uses a cached thread pool with the number of threads equal to the number of available processors.
     */
    public BegSmallZhV15Model(String modelPath, String tokenizerPath,PoolingMode poolingMode) {
        super(null);
        MODEL=  new OnnxBertBiEncoder(
                Path.of(modelPath),
                Path.of(tokenizerPath),
                poolingMode
        );
    }

    /**
     * Creates an instance of an {@code BegSmallZhV15Model}.
     *
     * @param executor The executor to use to parallelize the embedding process.
     */
    public BegSmallZhV15Model(Executor executor,String modelPath, String tokenizerPath,PoolingMode poolingMode) {
        super(ensureNotNull(executor, "executor"));
        MODEL=  new OnnxBertBiEncoder(
                Path.of(modelPath),
                Path.of(tokenizerPath),
                poolingMode
        );
    }

    @Override
    protected OnnxBertBiEncoder model() {
        return MODEL;
    }

    @Override
    protected Integer knownDimension() {
        return 512;
    }
}
